#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Python 执行 Shell 命令，将 STDOUT、STDERR 记录到指定日志（logging）
参考：https://stackoverflow.com/questions/803265/getting-realtime-output-using-subprocess/53323746#53323746
"""

import argparse
import asyncio
import logging
import shlex
from asyncio.subprocess import PIPE

LOG_LEVEL = logging.DEBUG
LOG_FORMAT = '%(asctime)s [%(levelname)s] %(message)s'
LOG_FORMATTER = logging.Formatter(LOG_FORMAT)
logging.basicConfig(level=LOG_LEVEL, format=LOG_FORMAT)

# def run(cmd, logfile):
#     with open(logfile, mode='bw', encoding='utf-8') as _file:
#         proc = Popen([cmd], bufsize=512, stdout=PIPE, stderr=PIPE, shell=True)
#         # proc.stderr is an io.TextIOWrapper file-like obj
#         for line in proc.stderr:
#             # print('current line:', line, flush=True)
#             _file.write(line)  # file open in binary mode
#             # _file.write(line.decode('utf8')) # for text mode
#             _file.flush()  # flush the content


async def _read_stream(stream, callback):
    while True:
        line = await stream.readline()
        if line:
            callback(line)
        else:
            break


async def run(command):
    process = await asyncio.create_subprocess_shell(
        *command, stdout=PIPE, stderr=PIPE
    )

    await asyncio.wait(
        [
            _read_stream(
                process.stdout,
                # lambda x: print('STDOUT: {}'.format(x.decode('UTF8'))),
                lambda x: logging.info(x.decode('UTF8').rstrip('\n')),
            ),
            _read_stream(
                process.stderr,
                # lambda x: print('STDERR: {}'.format(x.decode('UTF8'))),
                lambda x: logging.error(x.decode('UTF8').rstrip('\n')),
            ),
        ]
    )

    await process.wait()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', dest='command')
    parser.add_argument('-o', dest='logfile')
    args = parser.parse_args()

    if args.logfile:
        handler = logging.FileHandler(args.logfile, mode='w', encoding='utf-8')
        handler.setLevel(LOG_LEVEL)
        handler.setFormatter(LOG_FORMATTER)
        logging.root.addHandler(handler)

    async def worker():
        await run(shlex.split(args.command))

    loop = asyncio.get_event_loop()
    loop.set_debug(True)
    loop.run_until_complete(worker())


if __name__ == '__main__':
    main()
